import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import torch
import seaborn as sns


def calc_portfolio_metrics(
    portfolio_value: np.ndarray, risk_free_rate=0.0, benchmark=None
):
    portfolio_value = np.asarray(portfolio_value)

    # 普通收益率
    daily_return = portfolio_value[1:] / portfolio_value[:-1] - 1

    # 年化参数
    trading_num = 252 * 69
    mean_daily_return = np.mean(daily_return)
    std_daily_return = np.std(daily_return)

    ARR = (portfolio_value[-1] / portfolio_value[0]) ** (
        trading_num / len(portfolio_value)
    ) - 1
    AVol = std_daily_return * np.sqrt(trading_num)

    peak = np.maximum.accumulate(portfolio_value)
    MDD = np.max((peak - portfolio_value) / peak)

    ASR = (
        (mean_daily_return - risk_free_rate) / std_daily_return * np.sqrt(trading_num)
        if std_daily_return != 0
        else np.nan
    )

    # IR（相对于基准）
    if benchmark is not None:
        benchmark = np.asarray(benchmark)
        benchmark_return = benchmark[1:] / benchmark[:-1] - 1
        excess_ret = daily_return - benchmark_return
        IR = (
            np.mean(excess_ret) / np.std(excess_ret) * np.sqrt(trading_num)
            if np.std(excess_ret) != 0
            else np.nan
        )
    else:
        IR = np.nan

    return {"ARR": ARR, "AVol": AVol, "MDD": MDD, "ASR": ASR, "IR": IR}


def plot_portfolio_value(portfolio_value, benchmark=None):
    portfolio_value = np.array(portfolio_value, dtype=np.float32)
    portfolio_value = portfolio_value / portfolio_value[0]  # 标准化，起始为1

    plt.figure(figsize=(10, 5))
    plt.plot(portfolio_value, label="Agent Portfolio", linewidth=2)

    if benchmark is not None:
        benchmark = np.array(benchmark, dtype=np.float32)
        benchmark = benchmark / benchmark[0]  # 同样标准化
        plt.plot(benchmark, label="Benchmark", linestyle="--", linewidth=2)

    plt.title("Normalized Portfolio Value Over Time")
    plt.xlabel("Time Step")
    plt.ylabel("Normalized Portfolio Value")
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.show()


def plot_drawdown_curve(portfolio_value):
    peak = np.maximum.accumulate(portfolio_value)
    drawdown = (portfolio_value - peak) / peak

    plt.figure(figsize=(10, 3))
    plt.plot(drawdown, color="red", linewidth=2)
    plt.title("Drawdown Curve")
    plt.xlabel("Time Step")
    plt.ylabel("Drawdown (%)")
    plt.grid(True)
    plt.tight_layout()
    plt.show()


def plot_metric_bar_chart(metric_dict_list, strategy_names):
    metrics = ["ARR", "AVol", "MDD", "ASR", "IR"]
    metric_matrix = np.array([[m[key] for key in metrics] for m in metric_dict_list])

    x = np.arange(len(metrics))
    bar_width = 0.2

    plt.figure(figsize=(10, 5))
    for i, strategy_metrics in enumerate(metric_matrix):
        plt.bar(
            x + i * bar_width,
            strategy_metrics,
            width=bar_width,
            label=strategy_names[i],
        )

    plt.xticks(x + bar_width * (len(strategy_names) - 1) / 2, metrics)
    plt.title("Strategy Metric Comparison")
    plt.legend()
    plt.grid(axis="y")
    plt.tight_layout()
    plt.show()


def plot_price_and_actions(
    prices,
    actions,
    title="Price and Actions",
    save_dir=None,
    index=None,
    length=None,
    epoch=None,
):
    prices = np.array(prices)
    actions = np.array(
        [a.cpu().detach().numpy() if hasattr(a, "detach") else a for a in actions]
    )
    if index is not None:
        start = index * length
        end = min(start + length, len(prices))  # 防止越界
        prices = prices[start:end]

    # 裁剪长度一致
    min_len = min(len(prices), len(actions))
    prices = prices[:min_len]
    actions = actions[:min_len]

    n_assets = prices.shape[1]
    fig, axes = plt.subplots(n_assets, 1, figsize=(14, 3.5 * n_assets), sharex=True)

    if n_assets == 1:
        axes = [axes]  # 统一为 list

    for idx in range(n_assets):
        price_series = prices[:, idx]
        action_series = actions[:, idx]

        ax1 = axes[idx]
        ax1.set_title(f"{title} - Asset {idx}")
        ax1.plot(price_series, color="blue", label="Price")
        ax1.set_ylabel("Price", color="blue")
        ax1.tick_params(axis="y", labelcolor="blue")

        # Twin y-axis for action
        ax2 = ax1.twinx()
        ax2.plot(action_series, color="orange", label="Action", linestyle="--")
        ax2.set_ylabel("Action", color="orange")
        ax2.tick_params(axis="y", labelcolor="orange")
        ax2.axhline(0, color="gray", linestyle="--", linewidth=1)

        # 图例（左下角防重叠）
        ax1.legend(loc="lower left")
        ax2.legend(loc="lower right")

    plt.xlabel("Time Step")
    plt.tight_layout()

    if save_dir:
        os.makedirs(save_dir, exist_ok=True)  # 自动检查并创建目录
        filename = f"epoch_{epoch}_{start}_to_{end}.jpg"
        save_path = os.path.join(save_dir, filename)
        plt.savefig(save_path, dpi=300, bbox_inches="tight")
        plt.close()
    else:
        plt.show()


def show_win_rate(prices, actions):
    """
    计算每支期货的胜率（适配 actions 形状为 [time_steps, num_futures]）

    参数:
        prices (np.ndarray): 价格数组，形状 (time_steps+1, num_futures)
        actions (list/np.ndarray): 操作数组，形状 (time_steps, num_futures)，取值任意浮点数

    返回:
        dict: 每支期货的胜率 {期货索引: 胜率}
    """
    if isinstance(prices, torch.Tensor):
        prices = prices.cpu().numpy()
    else:
        prices = np.array(prices)

    if isinstance(actions, torch.Tensor):
        actions = actions.detach().cpu().numpy()
    elif isinstance(actions[0], torch.Tensor):
        actions = torch.stack([a.detach() for a in actions]).cpu().numpy()
    else:
        actions = np.array(actions)

    time_steps_plus_1, num_futures = prices.shape
    time_steps, num_futures_action = actions.shape

    assert num_futures == num_futures_action, "期货数量不一致"
    assert time_steps_plus_1 == time_steps + 1, "price长度应比action多1"

    win_rates = {}

    # 全局统计变量
    total_correct = 0
    total_count = 0
    up_correct = 0
    up_total = 0
    down_correct = 0
    down_total = 0

    for i in range(num_futures):
        price_series = prices[:, i]
        action_series = actions[:, i]

        price_diff = np.diff(price_series)
        action_signals = action_series

        mask = price_diff != 0
        valid_diff = price_diff[mask]
        valid_signals = action_signals[mask]

        # 胜率统计
        correct = np.sum(valid_diff * valid_signals > 0)
        total = len(valid_diff)
        win_rates[i] = correct / total if total > 0 else 0.0

        # 全局统计累计
        total_correct += correct
        total_count += total

        up_mask = valid_diff > 0
        up_correct += np.sum(valid_signals[up_mask] > 0)
        up_total += np.sum(up_mask)

        down_mask = valid_diff < 0
        down_correct += np.sum(valid_signals[down_mask] < 0)
        down_total += np.sum(down_mask)

    # 打印全局准确率
    acc = total_correct / total_count if total_count > 0 else 0.0
    up_acc = up_correct / up_total if up_total > 0 else 0.0
    down_acc = down_correct / down_total if down_total > 0 else 0.0

    print(
        f"[Validation] Total: {total_count}, Correct: {total_correct}, Accuracy: {acc:.4f}"
    )
    print(f"  ↑ Up:     {up_correct}/{up_total} = {up_acc:.4f}")
    print(f"  ↓ Down:   {down_correct}/{down_total} = {down_acc:.4f}")

    return win_rates


def show_win_rate_and_confusion_matrix(prices, actions):
    """
    计算每支期货的胜率和混淆矩阵（适配 actions 形状为 [time_steps, num_futures]）

    参数:
        prices (np.ndarray): 价格数组，形状 (time_steps+1, num_futures)
        actions (list/np.ndarray): 操作数组，形状 (time_steps, num_futures)，取值任意浮点数

    返回:
        win_rates (dict): 每支期货的胜率 {期货索引: 胜率}
        confusion_matrices (dict): 每支期货的混淆矩阵 {期货索引: 2x2 np.array}
    """
    if isinstance(prices, torch.Tensor):
        prices = prices.cpu().numpy()
    else:
        prices = np.array(prices)

    if isinstance(actions, torch.Tensor):
        actions = actions.detach().cpu().numpy()
    elif isinstance(actions[0], torch.Tensor):
        actions = torch.stack([a.detach() for a in actions]).cpu().numpy()
    else:
        actions = np.array(actions)

    time_steps_plus_1, num_futures = prices.shape
    time_steps, num_futures_action = actions.shape

    assert num_futures == num_futures_action, "期货数量不一致"
    assert time_steps_plus_1 == time_steps + 1, "price 长度应比 action 多 1"

    win_rates = {}
    confusion_matrices = {}

    for i in range(num_futures):
        price_series = prices[:, i]  # (time_steps + 1,)
        action_series = actions[:, i]  # (time_steps,)

        price_diff = np.diff(price_series)  # (time_steps,)
        action_signals = action_series  # (time_steps,)

        # 过滤掉价格未变化的情况
        mask = price_diff != 0
        valid_price_diff = price_diff[mask]
        valid_action_signals = action_signals[mask]

        correct = np.sum(valid_price_diff * valid_action_signals > 0)
        total = len(valid_price_diff)

        win_rates[i] = correct / total if total > 0 else 0.0

        # 混淆矩阵构建
        actual_up = valid_price_diff > 0
        actual_down = valid_price_diff < 0
        predicted_up = valid_action_signals > 0
        predicted_down = valid_action_signals < 0

        tp = np.sum(actual_up & predicted_up)
        fp = np.sum(actual_down & predicted_up)
        fn = np.sum(actual_up & predicted_down)
        tn = np.sum(actual_down & predicted_down)

        confusion_matrix = np.array([[tp, fp], [fn, tn]])
        confusion_matrices[i] = confusion_matrix

    # 可视化
    for i in range(num_futures):
        plt.figure(figsize=(6, 4))
        sns.heatmap(
            confusion_matrices[i],
            annot=True,
            fmt="d",
            cmap="Blues",
            xticklabels=["Predicted Down", "Predicted Up"],
            yticklabels=["Actual Down", "Actual Up"],
        )
        plt.title(f"Confusion Matrix for Future {i}")
        plt.xlabel("Predicted")
        plt.ylabel("Actual")
        plt.tight_layout()
        plt.show()

    return win_rates, confusion_matrices


if __name__ == "__main__":
    np.random.seed(42)
    days = 100
    simulated_returns = np.random.normal(loc=0.001, scale=0.01, size=days)
    portfolio_value = [100]
    for r in simulated_returns:
        portfolio_value.append(portfolio_value[-1] * (1 + r))
    portfolio_value = np.array(portfolio_value)
    benchmark = 100 * np.cumprod(
        1 + np.random.normal(loc=0.0008, scale=0.009, size=days + 1)
    )
    metrics = calc_portfolio_metrics(portfolio_value, benchmark=benchmark)
    for k, v in metrics.items():
        print(f"{k}: {v:.4f}")
    metrics_agent = calc_portfolio_metrics(portfolio_value, benchmark=benchmark)
    metrics_random = calc_portfolio_metrics(
        portfolio_value * np.random.uniform(0.95, 1.05, size=len(portfolio_value))
    )
    plot_portfolio_value(portfolio_value, benchmark)
    # plot_drawdown_curve(portfolio_value)
    plot_metric_bar_chart(
        metric_dict_list=[metrics_agent, metrics_random],
        strategy_names=["RL Agent", "Random Baseline"],
    )
    # 构造假数据
    T = 100
    N = 6
    np.random.seed(0)

    # 假设价格在不断上涨
    fake_prices = np.cumsum(np.random.randn(T, N) * 2 + 1, axis=0) + 100

    # 假设 agent 做多/做空有点“跟随趋势”
    fake_actions = np.tanh(np.gradient(fake_prices, axis=0))

    # 可视化第0个资产的交易行为
    plot_price_and_actions(
        fake_prices, fake_actions, asset_index=0, title="Price vs Action (Asset 0)"
    )
